package ssh

import (
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

type ChMsgDirectTcpip struct {
	DestAddr   string
	DestPort   uint32
	OriginAddr string
	OriginPort uint32
}

type DialTcpHandler struct {
	CheckAuth   func(msg ChMsgDirectTcpip) error
	WatchConfig WatchConfigFunc
}

func (handler *DialTcpHandler) Handle(newCh ssh.NewChannel) {
	if newCh.ChannelType() != "direct-tcpip" {
		log.Panicf("ssh: unknown channel type %v", newCh.ChannelType())
	}
	var msg ChMsgDirectTcpip
	if err := ssh.Unmarshal(newCh.ExtraData(), &msg); err != nil {
		newCh.Reject(ssh.ConnectionFailed, fmt.Sprintf("could not parse %v payload: %v", newCh.ChannelType(), err.Error()))
		return
	}
	if err := handler.CheckAuth(msg); err != nil {
		newCh.Reject(ssh.Prohibited, err.Error())
		return
	}
	addr := JoinHostPort(msg.DestAddr, msg.DestPort)
	c, err := net.Dial("tcp", addr)
	if err != nil {
		newCh.Reject(ssh.ConnectionFailed, fmt.Sprintf("could not dial to %v: %v", addr, err.Error()))
		return
	}
	defer c.Close()
	conn := c.(*net.TCPConn)
	ch, reqs, err := newCh.Accept()
	if err != nil {
		return
	}
	defer ch.Close()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		watcher := handler.WatchConfig()
		defer watcher.Unwatch()
	Loop:
		for {
			select {
			case <-ctx.Done():
				break Loop
			case _, ok := <-watcher.OnChange:
				if !ok {
					break Loop
				}
			}
			if err := handler.CheckAuth(msg); err != nil {
				ch.Close()
			}
		}
	}()
	go ssh.DiscardRequests(reqs)
	copyWg := &sync.WaitGroup{}
	defer copyWg.Wait()
	copyWg.Add(2)
	go func() {
		defer copyWg.Done()
		io.Copy(conn, ch)
		conn.CloseWrite()
		// ch.CloseRead()
	}()
	go func() {
		defer copyWg.Done()
		io.Copy(ch, conn)
		ch.CloseWrite()
		conn.CloseRead()
	}()
}

type ReqMsgTcpipForward struct {
	Addr string
	Port uint32
}

type ReplyMsgTcpipForward struct {
	Port uint32
}

type ChMsgForwardedTcpip struct {
	Addr       string
	Port       uint32
	OriginAddr string
	OriginPort uint32
}

type ReqMsgCancelTcpipForward struct {
	Addr string
	Port uint32
}

type ListenTcpHandler struct {
	SrvConn     *ssh.ServerConn
	CheckAuth   func(msg ReqMsgTcpipForward) error
	WatchConfig WatchConfigFunc
	cancels     sync.Map
}

func (handler *ListenTcpHandler) Finalize() {
	handler.cancels.Range(func(key, value any) bool {
		value.(context.CancelFunc)()
		return true
	})
}

func (handler *ListenTcpHandler) forward(addr string) (net.Listener, error) {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				break
			}
			go func() {
				defer c.Close()
				conn := c.(*net.TCPConn)
				lAddr := conn.LocalAddr().(*net.TCPAddr)
				rAddr := conn.RemoteAddr().(*net.TCPAddr)
				ch, reqs, err := handler.SrvConn.OpenChannel("forwarded-tcpip", ssh.Marshal(&ChMsgForwardedTcpip{
					Addr:       lAddr.IP.String(),
					Port:       uint32(lAddr.Port),
					OriginAddr: rAddr.IP.String(),
					OriginPort: uint32(rAddr.Port),
				}))
				if err != nil {
					return
				}
				defer ch.Close()
				go ssh.DiscardRequests(reqs)
				copyWg := &sync.WaitGroup{}
				defer copyWg.Wait()
				copyWg.Add(2)
				go func() {
					defer copyWg.Done()
					io.Copy(conn, ch)
					conn.CloseWrite()
					// ch.CloseRead()
				}()
				go func() {
					defer copyWg.Done()
					io.Copy(ch, conn)
					ch.CloseWrite()
					conn.CloseRead()
				}()
			}()
		}
	}()
	return ln, nil
}

func (handler *ListenTcpHandler) Handle(req *ssh.Request) {
	Reply(req, func() (bool, []byte) {
		if req.Type != "tcpip-forward" {
			log.Panicf("ssh: unknown channel type %v", req.Type)
		}
		var msg ReqMsgTcpipForward
		if err := ssh.Unmarshal(req.Payload, &msg); err != nil {
			return false, nil
		}
		if err := handler.CheckAuth(msg); err != nil {
			return false, nil
		}
		ln, err := handler.forward(JoinHostPort(msg.Addr, msg.Port))
		if err != nil {
			return false, nil
		}
		addr := ln.Addr().(*net.TCPAddr)
		ctx, cancel := context.WithCancel(context.Background())
		if previous, loaded := handler.cancels.Swap(addr, cancel); loaded {
			previous.(context.CancelFunc)()
		}
		watchWg := &sync.WaitGroup{}
		watchWg.Add(1)
		go func() {
			<-ctx.Done()
			watchWg.Wait()
			ln.Close()
		}()
		go func() {
			defer watchWg.Done()
			watcher := handler.WatchConfig()
			defer watcher.Unwatch()
			ticker := time.NewTicker(time.Minute)
			defer ticker.Stop()
		Loop:
			for {
				select {
				case <-ctx.Done():
					break Loop
				case _, ok := <-watcher.OnChange:
					if !ok {
						break Loop
					}
					ticker.Reset(time.Minute)
				case <-ticker.C:
				}
				if err := handler.CheckAuth(msg); err != nil {
					if ln != nil {
						ln.Close()
						ln = nil
						handler.SrvConn.SendRequest("cancel-tcpip-forward", true, ssh.Marshal(&ReqMsgCancelTcpipForward{
							Addr: addr.IP.String(),
							Port: uint32(addr.Port),
						}))
					}
				} else {
					if ln == nil {
						ln, _ = handler.forward(addr.String())
					}
				}
			}
		}()
		return true, ssh.Marshal(&ReplyMsgTcpipForward{
			Port: uint32(ln.Addr().(*net.TCPAddr).Port),
		})
	})
}

func (handler *ListenTcpHandler) HandleClose(req *ssh.Request) {
	Reply(req, func() (bool, []byte) {
		if req.Type != "cancel-tcpip-forward" {
			log.Panicf("ssh: unknown channel type %v", req.Type)
		}
		var msg ReqMsgCancelTcpipForward
		if err := ssh.Unmarshal(req.Payload, &msg); err != nil {
			return false, nil
		}
		value, loaded := handler.cancels.LoadAndDelete(JoinHostPort(msg.Addr, msg.Port))
		if !loaded {
			return false, nil
		}
		value.(context.CancelFunc)()
		return true, nil
	})
}
